// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/expression/fuzzer_runner.h>

#include <pollux/testing/expression/expression_fuzzer.h>
#include <pollux/testing/flags.h>

namespace kumo::pollux::fuzzer {
    namespace {
        VectorFuzzer::Options getVectorFuzzerOptions() {
            VectorFuzzer::Options opts;
            opts.vectorSize = turbo::get_flag(FLAGS_batch_size);
            opts.stringVariableLength = true;
            opts.stringLength = 100;
            opts.nullRatio = turbo::get_flag(FLAGS_null_ratio);
            opts.timestampPrecision =
                    VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
            return opts;
        }

        ExpressionFuzzer::Options getExpressionFuzzerOptions(
            const std::unordered_set<std::string> &skipFunctions,
            const std::unordered_map<std::string, std::shared_ptr<ExprTransformer> > &
            exprTransformers,
            std::shared_ptr<ReferenceQueryRunner> referenceQueryRunner) {
            ExpressionFuzzer::Options opts;
            opts.maxLevelOfNesting = turbo::get_flag(FLAGS_pollux_fuzzer_max_level_of_nesting);
            opts.maxNumVarArgs = turbo::get_flag(FLAGS_max_num_varargs);
            opts.enableVariadicSignatures = turbo::get_flag(FLAGS_enable_variadic_signatures);
            opts.enableDereference = turbo::get_flag(FLAGS_enable_dereference);
            opts.enableComplexTypes = turbo::get_flag(FLAGS_pollux_fuzzer_enable_complex_types);
            opts.enableDecimalType = turbo::get_flag(FLAGS_pollux_fuzzer_enable_decimal_type);
            opts.enableColumnReuse = turbo::get_flag(FLAGS_pollux_fuzzer_enable_column_reuse);
            opts.enableExpressionReuse = turbo::get_flag(FLAGS_pollux_fuzzer_enable_expression_reuse);
            opts.functionTickets = turbo::get_flag(FLAGS_assign_function_tickets);
            opts.nullRatio = turbo::get_flag(FLAGS_null_ratio);
            opts.specialForms = turbo::get_flag(FLAGS_special_forms);
            opts.useOnlyFunctions = turbo::get_flag(FLAGS_only);
            opts.skipFunctions = skipFunctions;
            opts.referenceQueryRunner = referenceQueryRunner;
            opts.exprTransformers = exprTransformers;
            return opts;
        }

        ExpressionFuzzerVerifier::Options getExpressionFuzzerVerifierOptions(
            const std::unordered_set<std::string> &skipFunctions,
            const std::unordered_map<std::string, std::shared_ptr<ExprTransformer> > &
            exprTransformers,
            const std::unordered_map<std::string, std::string> &queryConfigs,
            std::shared_ptr<ReferenceQueryRunner> referenceQueryRunner) {
            ExpressionFuzzerVerifier::Options opts;
            opts.steps = turbo::get_flag(FLAGS_steps);
            opts.durationSeconds = turbo::get_flag(FLAGS_duration_sec);
            opts.batchSize = turbo::get_flag(FLAGS_batch_size);
            opts.retryWithTry = turbo::get_flag(FLAGS_retry_with_try);
            opts.findMinimalSubexpression = turbo::get_flag(FLAGS_find_minimal_subexpression);
            opts.disableConstantFolding = turbo::get_flag(FLAGS_disable_constant_folding);
            opts.reproPersistPath = turbo::get_flag(FLAGS_repro_persist_path);
            opts.persistAndRunOnce = turbo::get_flag(FLAGS_persist_and_run_once);
            opts.lazyVectorGenerationRatio = turbo::get_flag(FLAGS_lazy_vector_generation_ratio);
            opts.commonDictionaryWrapRatio =
                    turbo::get_flag(FLAGS_common_dictionary_wraps_generation_ratio);
            opts.maxExpressionTreesPerStep = turbo::get_flag(FLAGS_max_expression_trees_per_step);
            opts.vectorFuzzerOptions = getVectorFuzzerOptions();
            opts.expressionFuzzerOptions = getExpressionFuzzerOptions(
                skipFunctions, exprTransformers, referenceQueryRunner);
            opts.queryConfigs = queryConfigs;
            return opts;
        }
    } // namespace

    // static
    int FuzzerRunner::run(
        size_t seed,
        const std::unordered_set<std::string> &skipFunctions,
        const std::unordered_map<std::string, std::shared_ptr<ExprTransformer> > &
        exprTransformers,
        const std::unordered_map<std::string, std::string> &queryConfigs,
        const std::unordered_map<std::string, std::shared_ptr<ArgTypesGenerator> > &
        argTypesGenerators,
        const std::unordered_map<std::string, std::shared_ptr<ArgValuesGenerator> > &
        argValuesGenerators,
        std::shared_ptr<ReferenceQueryRunner> referenceQueryRunner,
        const std::shared_ptr<SpecialFormSignatureGenerator> &
        specialFormSignatureGenerator) {
        runFromGtest(
            seed,
            skipFunctions,
            exprTransformers,
            queryConfigs,
            argTypesGenerators,
            argValuesGenerators,
            referenceQueryRunner,
            specialFormSignatureGenerator);
        return RUN_ALL_TESTS();
    }

    // static
    void FuzzerRunner::runFromGtest(
        size_t seed,
        const std::unordered_set<std::string> &skipFunctions,
        const std::unordered_map<std::string, std::shared_ptr<ExprTransformer> > &
        exprTransformers,
        const std::unordered_map<std::string, std::string> &queryConfigs,
        const std::unordered_map<std::string, std::shared_ptr<ArgTypesGenerator> > &
        argTypesGenerators,
        const std::unordered_map<std::string, std::shared_ptr<ArgValuesGenerator> > &
        argValuesGenerators,
        std::shared_ptr<ReferenceQueryRunner> referenceQueryRunner,
        const std::shared_ptr<SpecialFormSignatureGenerator> &
        specialFormSignatureGenerator) {
        if (!memory::MemoryManager::testInstance()) {
            memory::MemoryManager::testingSetInstance({});
        }
        auto signatures = kumo::pollux::get_function_signatures();
        const auto options = getExpressionFuzzerVerifierOptions(
            skipFunctions, exprTransformers, queryConfigs, referenceQueryRunner);
        // Insert generated signatures of special forms into the signature map.
        specialFormSignatureGenerator->appendSpecialForms(
            signatures, options.expressionFuzzerOptions.specialForms);
        ExpressionFuzzerVerifier(
                    signatures, seed, options, argTypesGenerators, argValuesGenerators)
                .go();
    }
} // namespace kumo::pollux::fuzzer
